package com.youngjun.common.security.filter;

import com.youngjun.common.security.service.RequestMapService;
import com.youngjun.common.security.validator.X509CertificateValidator;
import com.youngjun.common.security.constant.PKIContat;
import com.youngjun.common.util.SpringUtils;
import com.youngjun.common.util.PKIUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.security.authentication.*;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.AuthenticationException;
import org.springframework.security.core.context.SecurityContextHolder;
import org.springframework.security.web.authentication.WebAuthenticationDetailsSource;
import org.springframework.security.web.authentication.preauth.PreAuthenticatedAuthenticationToken;
import org.springframework.security.web.authentication.preauth.x509.SubjectDnX509PrincipalExtractor;
import org.springframework.security.web.authentication.preauth.x509.X509PrincipalExtractor;

import javax.servlet.FilterChain;
import javax.servlet.ServletException;
import javax.servlet.ServletRequest;
import javax.servlet.ServletResponse;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.servlet.http.HttpSession;
import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.cert.CertificateException;
import java.security.cert.X509Certificate;

public class X509AuthenticationFilter extends org.springframework.security.web.authentication.preauth.x509.X509AuthenticationFilter {

    private final static Logger logger = LoggerFactory.getLogger(X509AuthenticationFilter.class);

    private X509PrincipalExtractor principalExtractor = new SubjectDnX509PrincipalExtractor();
    private AuthenticationDetailsSource<HttpServletRequest, ?> authenticationDetailsSource = new WebAuthenticationDetailsSource();
    private AuthenticationManager authenticationManager = null;
    private boolean checkForPrincipalChanges;
    private boolean invalidateSessionOnPrincipalChange = true;
    private X509CertificateValidator x509CertificateValidator;

    private final String SSL_CLIENT_CERT_HEADER = "x-ssl-client-cert";

    public void setX509CertificateValidator(X509CertificateValidator x509CertificateValidator) {
        this.x509CertificateValidator = x509CertificateValidator;
    }
    public void setCheckForPrincipalChanges(boolean checkForPrincipalChanges) {
        this.checkForPrincipalChanges = checkForPrincipalChanges;
        super.setCheckForPrincipalChanges(checkForPrincipalChanges);
    }
    public void setInvalidateSessionOnPrincipalChange(
            boolean invalidateSessionOnPrincipalChange) {
        this.invalidateSessionOnPrincipalChange = invalidateSessionOnPrincipalChange;
        super.setInvalidateSessionOnPrincipalChange(invalidateSessionOnPrincipalChange);
    }
    public void setAuthenticationManager(AuthenticationManager authenticationManager) {
        this.authenticationManager = authenticationManager;
        super.setAuthenticationManager(authenticationManager);
    }

    protected Object getPreAuthenticatedPrincipal(HttpServletRequest request) {
        X509Certificate cert = extractClientCertificate(request);

        if (cert == null) {
            return null;
        }

        return principalExtractor.extractPrincipal(cert);
    }

    protected Object getPreAuthenticatedCredentials(HttpServletRequest request) {
        return extractClientCertificate(request);
    }

    private X509Certificate extractClientCertificate(HttpServletRequest request) {
        String headerValue = request.getHeader(SSL_CLIENT_CERT_HEADER);
        if (headerValue != null) {
            X509Certificate cert = null;
            try {
                cert = PKIUtils.loadCertificateFromPEM(PKIContat.DEFAULT_CERTIFICATE_TYPE, new ByteArrayInputStream(headerValue.getBytes(StandardCharsets.UTF_8)));
                if (logger.isDebugEnabled()) {
                    logger.debug("X.509 client authentication certificate:" + cert);
                }
                return cert;
            } catch (CertificateException | IOException e) {
                e.printStackTrace();
            }
        }
        if (logger.isDebugEnabled()) {
            logger.debug("No client certificate found in request.");
        }
        return null;
    }

    public void setPrincipalExtractor(X509PrincipalExtractor principalExtractor) {
        this.principalExtractor = principalExtractor;
    }

    public void doFilter(ServletRequest request, ServletResponse response,
                         FilterChain chain) throws IOException, ServletException {

        if (logger.isDebugEnabled()) {
            logger.debug("Checking secure context token: "
                    + SecurityContextHolder.getContext().getAuthentication());
        }

        if (requiresAuthentication((HttpServletRequest) request)) {
            doAuthenticate((HttpServletRequest) request, (HttpServletResponse) response);
        }

        chain.doFilter(request, response);
    }

    /**
     * Do the actual authentication for a pre-authenticated user.
     */
    private void doAuthenticate(HttpServletRequest request, HttpServletResponse response) throws IOException, ServletException {
        Authentication authResult;

        Object principal = getPreAuthenticatedPrincipal(request);
        Object credentials = getPreAuthenticatedCredentials(request);

        if (principal == null) {
            if (logger.isDebugEnabled()) {
                logger.debug("No pre-authenticated principal found in request");
            }

            return;
        }

        if (credentials == null){
            if (logger.isDebugEnabled()) {
                logger.debug("No pre-authenticated credentials found in request");
            }

            return;
        }

        if (logger.isDebugEnabled()) {
            logger.debug("preAuthenticatedPrincipal = " + principal
                    + ", trying to authenticate");
        }

        try {
            x509CertificateValidator.validate((X509Certificate) credentials);
            PreAuthenticatedAuthenticationToken authRequest = new PreAuthenticatedAuthenticationToken(
                    principal, credentials);
            authRequest.setDetails(authenticationDetailsSource.buildDetails(request));
            authResult = authenticationManager.authenticate(authRequest);
            successfulAuthentication(request, response, authResult);
        }catch (AuthenticationException failed) {
            unsuccessfulAuthentication(request, response, failed);
            RequestMapService requestMapService = SpringUtils.getBean(RequestMapService.class);
            if (requestMapService.mustX509Authentication(request.getServletPath(),request.getMethod())) {
                throw failed;
            }
        }
    }


    private boolean requiresAuthentication(HttpServletRequest request) {
        Authentication currentUser = SecurityContextHolder.getContext()
                .getAuthentication();

        if (currentUser == null) {
            return true;
        }

        if (!checkForPrincipalChanges) {
            return false;
        }

        if(!principalChanged(request, currentUser)) {
            return false;
        }

        logger.debug("Pre-authenticated principal has changed and will be reauthenticated");

        if (invalidateSessionOnPrincipalChange) {
            SecurityContextHolder.clearContext();

            HttpSession session = request.getSession(false);

            if (session != null) {
                logger.debug("Invalidating existing session");
                session.invalidate();
                request.getSession();
            }
        }

        return true;
    }

}
